"""用于仿真生成标注文件、表格文件的上层脚本"""

from gt_gen import ChangE1_5m, SegmentGT
from catalog import XuLihengCatalog, ConicPointCatalog
import os
import yaml
import cv2
import numpy as np
import tqdm
import argparse


def arg_parse():
    arg_parser = argparse.ArgumentParser()
    arg_parser.add_argument(
        "-c",
        "--collect_dir",
        type=str,
        help="The directory of the collected data",
    )
    arg_parser.add_argument(
        "-l",
        "--label_dir",
        type=str,
        default="/disk527/sdb1/a804_cbf/datasets/lunar_crater",
        help="The directory of the label data",
    )
    arg_parser.add_argument(
        "-o",
        "--catalog_dir",
        type=str,
        default="/disk527/sdb1/a804_cbf/catalog",
        help="The directory of the catalog",
    )
    arg_parser.add_argument(
        "-d",
        "--max_dist",
        type=float,
        default=240,
        help="The maximum distance between two invariant points",
    )
    arg_parser.add_argument("--cache", action="store_true", help="Whether to use cache")
    arg_parser.add_argument(
        "--image", action="store_true", help="Whether to use image as background"
    )
    arg_parser.add_argument(
        "--catalog",
        action="store_true",
        default=False,
        help="Whether to generate catalog",
    )
    arg_parser.add_argument(
        "--contour",
        action="store_true",
        help="Whether to generate contour",
    )
    arg_parser.add_argument(
        "--segment",
        action="store_true",
        help="Whether to generate segment",
    )
    arg_parser.add_argument(
        "-m",
        "--method",
        choices=["conic", "xuliheng"],
        help="The method to generate catalog",
    )
    return arg_parser.parse_args()


LATITUDE_RANGE = 45.2410 - 42.8613  # deg
LONGITUDE_RANGE = -51.6879 - (-52.0974)  # deg
X_PIXEL_RANGE = 48131
Y_PIXEL_RANGE = 8264
BASE_BATCH = 1000
num_batches = X_PIXEL_RANGE // BASE_BATCH
num_samples = Y_PIXEL_RANGE // BASE_BATCH
rate = 2
scale_factor = 0.2
modifier = 1.5 * rate / scale_factor


def gt_contour(sub_dir, collected_dir, cache_dir):
    with open(os.path.join(collected_dir, "config.yaml"), "r") as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    row = config["world"]["row"]
    col = config["world"]["col"]
    num_batches = X_PIXEL_RANGE // BASE_BATCH
    num_samples = Y_PIXEL_RANGE // BASE_BATCH
    delta_theta = LATITUDE_RANGE / num_batches
    delta_phi = LONGITUDE_RANGE / num_samples
    origin_point = (
        45.2410 - delta_theta * row[len(row) // 2],
        -52.0974 + delta_phi * col[len(col) // 2],
    )

    gt_gen = ChangE1_5m(
        sub_dir,
        origin_point,
        config["camera"]["P"],
        (config["camera"]["height"], config["camera"]["width"]),
        rate,
        cache_dir=cache_dir,
        modifier=modifier,
        lat=(45.2410, 42.8613),
        lon=(-52.0974, -51.6879),
        pixel_range=(X_PIXEL_RANGE, Y_PIXEL_RANGE),
        row=row,
        col=col,
    )
    image_names = list(
        map(
            lambda x: float(x.removesuffix(".png")),
            os.listdir(f"{collected_dir}/images"),
        )
    )
    image_names.sort()
    it = 0
    image_name = image_names.pop(0)
    with open(f"{collected_dir}/pose.csv", "r") as f:
        f.readline()  # 去除标题
        with open(f"{collected_dir}/gt_labels.txt", "w") as writer:
            end_of_file = False
            for line in tqdm.tqdm(
                f, desc="Generating GT images", total=len(image_names)
            ):
                # 某个时刻的姿态
                data = line.split(",")
                time = float(data[0])
                if image_name - time > 2e-4:
                    continue
                else:
                    time = f"{image_name:.6f}"
                    if len(image_names) != 0:
                        image_name = image_names.pop(0)
                    else:
                        end_of_file = True

                pose = list(float(x) for x in data[1:4])
                ori = list(float(x) for x in data[4:])
                img = cv2.imread(f"{collected_dir}/images/{time}.png")

                if np.std(img) < 2:
                    continue
                flag = gt_gen(
                    time, pose, ori, collected_dir, img if args.image else None
                )
                if flag:
                    print(f"{time}", file=writer, flush=True)
                    it += 1
                if end_of_file:
                    break
            print("=" * 20 + f"\ntotal labels: {it}\n" + "=" * 20)


def gt_segment(collect_dir, is_bg: bool):
    for collect_dir in collect_dir.split(","):
        with open(os.path.join(collect_dir, "config.yaml"), "r") as f:
            config = yaml.load(f.read(), Loader=yaml.FullLoader)
            size = (config["camera"]["height"], config["camera"]["width"])
        with SegmentGT(size, collect_dir) as gt_gen:
            gt_gen(is_bg)


def to_catalog(collect_dir, label_dir, catalog_dir, is_cache, max_dist, method="conic"):
    with open(os.path.join(collect_dir, "config.yaml"), "r") as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    row = config["world"]["row"]
    col = config["world"]["col"]
    # row = [42, 43, 44, 45, 46, 47]
    # col = [0, 1, 2, 3, 4, 5, 6, 7]
    delta_theta = LATITUDE_RANGE / num_batches
    delta_phi = LONGITUDE_RANGE / num_samples
    # 原点会影响不变量的计算，但是算好后的不变量，在不同原点对应的投影平面上都应该是相等的
    origin_point = [
        45.2410 - delta_theta * row[len(row) // 2],
        -52.0974 + delta_phi * col[len(col) // 2],
    ]
    if method == "xuliheng":
        catalog_method = XuLihengCatalog
    elif method == "conic":
        catalog_method = ConicPointCatalog
    catalog = catalog_method(
        label_dir,
        catalog_dir,
        row,
        col,
        origin_point,
        base_batch=BASE_BATCH,
        rate=rate,
        cache_dir="." if is_cache else None,
        modifier=modifier,
        lat=(45.2410, 42.8613),
        lon=(-52.0974, -51.6879),
        pixel_range=(X_PIXEL_RANGE, Y_PIXEL_RANGE),
    )
    catalog.to_catalog()
    catalog.to_invariant(max_dist=max_dist)


if __name__ == "__main__":
    args = arg_parse()
    for collect_dir in args.collect_dir.split(","):
        if args.catalog:
            to_catalog(
                collect_dir,
                args.label_dir,
                args.catalog_dir,
                args.cache,
                args.max_dist,
                args.method,
            )
        if args.contour:
            gt_contour(args.label_dir, collect_dir, "." if args.cache else None)
        if args.segment:
            gt_segment(collect_dir, args.image)
